# import json
# from arg_parse import load_parse
# args = load_parse()
# test_batch_size = args.test_batch_size
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '1.'
# with open(f'save/{args.title}/config.json', 'r') as f:
#     args.__dict__ = json.load(f)
# from data import load_data_hym
# import jax
# from mtpp import MTPP
# import equinox as eqx
# from eval import eval_nll, eval_one_step_predict

# # Data loading
# if args.dataname in ['stackoverflow', 'retweet', 'mimic', 'taobao', 'taxi', 'earthquake']:
#     test_data = load_data_hym(args.dataname, 'test')
#     train_data = load_data_hym(args.dataname, 'train')
#     val_data = load_data_hym(args.dataname, 'dev')
# test_loader = test_data.get_dataloader(test_batch_size)

# # Process horizon value for CTPP
# if args.method == 'ctpp':
#     channels = args.horizon.split(';')
#     horizon = [c.split(',') for c in channels]
#     horizon = jax.tree.map(lambda x: float(x), horizon)
# else:
#     horizon = None

# key = jax.random.PRNGKey(69)
# init_key, test_key = jax.random.split(key)

# # Get the maximum sequence length
# max_len = test_data.max_length()

# dt_max = train_data.get_dtmax()
# dt_mean = train_data.dt_stat()[1]

# temp = test_data.dt_stat()[-1]
# if args.max_dt > 0:
#     scale = dt_max / args.max_dt
# else:
#     scale = dt_mean


# model = MTPP(args.method, args.hdim, train_data.num_types(), args.hdim, init_key, {
#     'num_components': args.components, 'nlayers': args.layers, 'nhead': args.nhead, 'horizon': horizon, 
#     'omega': args.omega, 'siren_layers': args.siren_layers, 'dec_type': args.dec_type, 'attn_type': args.attn_type, 
#     'reg': args.reg
# })
# model = eqx.tree_deserialise_leaves(f'save/{args.title}/model', model)
# for i in range(test_data.times.shape[0]):
#     # out = model(test_data.times[i], test_data.marks[i], test_data.mask[i], jax.random.PRNGKey(1234))
#     out = model.sequence_one_step_predict(test_data.times[i], test_data.marks[i], test_data.mask[i], jax.random.PRNGKey(1234), 5/dt_mean)

# import json
# from arg_parse import load_parse
# args = load_parse()
# test_batch_size = args.test_batch_size
# import os
# os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
# os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '1.'
# with open(f'save/{args.title}/config.json', 'r') as f:
#     args.__dict__ = json.load(f)

from mtpp import MTPP
import jax
import jax.numpy as jnp
import equinox as eqx


model = MTPP('njdtpp2', 64, 73, 64, jax.random.PRNGKey(0), {'nlayers': 2, 'nhead': 3, 'num_components': 64, 'reg': 0., 'num_steps': 50})

# # model = eqx.filter_eval_shape(MTPP, args.method, args.hdim, 73, args.hdim, jax.random.PRNGKey(69), {
# #     'num_components': args.components, 'nlayers': args.layers, 'nhead': args.nhead, 
# #     'omega': args.omega, 'siren_layers': args.siren_layers, 'dec_type': args.dec_type, 'attn_type': args.attn_type, 
# #     'reg': args.reg
# # })
# # model = eqx.tree_deserialise_leaves(f'save/{args.title}/model', model)
# model = MTPP(args.method, args.hdim, 73, args.hdim, jax.random.PRNGKey(69), {
#     'num_components': args.components, 'nlayers': args.layers, 'nhead': args.nhead, 
#     'omega': args.omega, 'siren_layers': args.siren_layers, 'dec_type': args.dec_type, 'attn_type': args.attn_type, 
#     'reg': args.reg
# })
# model = eqx.tree_deserialise_leaves(f'save/{args.title}/model', model)


ts = jax.random.uniform(jax.random.PRNGKey(69), (16, 20)).cumsum(1)
marks = jax.random.randint(jax.random.PRNGKey(3234), (16, 20), 0, 73)
mask = jnp.ones_like(ts, dtype=bool)
out = jax.vmap(model, (0, 0, 0, 0))(ts, marks, mask, jax.random.split(jax.random.PRNGKey(1234), ts.shape[0]))
print()

# mask = jnp.ones_like(ts, dtype=bool)
# key = jax.random.PRNGKey(1234)
# # while True:
# #     cur_key, key = jax.random.split(key)
# #     Z0 = jax.random.uniform(cur_key, (73, args.hdim))
# #     Efdt, mark_predict = model.model._predict(Z0, 10.)
# #     assert Efdt >= 0, "Predicted time deltas must be non-negative."
# # print()
# for i in range(4):
#     out = model(ts[i], marks[i], mask[i], jax.random.PRNGKey(1234))
# #     # t_i = ts[i]
# #     # m_i = marks[i]
# #     # mask_i = mask[i]
# #     predict_tuple, real_tuple, mask_out = model.sequence_one_step_predict(ts[i], marks[i], mask[i], jax.random.PRNGKey(1234), 10.0)
# # out = jax.vmap(model, (0, 0, 0, 0))(ts, marks, mask, jax.random.split(jax.random.PRNGKey(1234), ts.shape[0]))
# # predict_tuple, real_tuple, mask = jax.vmap(model.sequence_one_step_predict, (0, 0, 0, 0, None))(ts, marks, mask, jax.random.split(jax.random.PRNGKey(1234), ts.shape[0]), 10.0)
# # assert jnp.all(predict_tuple[0] >= 0), "Predicted time deltas must be non-negative."
# # print()